import math
from collections import Counter
import jieba
import json
from loguru import logger
from tqdm import tqdm
import heapq


class BM25Searcher:
    def __init__(self, docs, k1=1.5, b=0.75):
        """
        BM25算法的构造器
        :param docs: 未分词的文档
        :param k1: BM25算法中的调节参数k1
        :param b: BM25算法中的调节参数b
        """
        self.docs = docs
        self.k1 = k1
        self.b = b
        self.doc_len = []  # 计算每个文档的长度
        self.avgdl = None  # 计算所有文档的平均长度
        self.doc_freqs = []  # 存储每个文档的词频
        self.idf = {}  # 存储每个词的逆文档频率
        self.initialize()

    def initialize(self):
        """
        初始化方法，计算所有词的逆文档频率
        """
        df = {}  # 用于存储每个词在多少不同文档中出现
        num_tokens = 0
        logger.info('统计词频...')
        for doc in tqdm(self.docs):
            # 为每个文档创建一个词频统计
            doc_seg = jieba.lcut_for_search(doc['content'])
            num_tokens += len(doc_seg)
            self.doc_len.append(len(doc_seg))
            self.doc_freqs.append(Counter(doc_seg))
            # 更新df值
            for word in set(doc_seg):
                df[word] = df.get(word, 0) + 1
        # 计算每个词的IDF值
        for word, freq in df.items():
            self.idf[word] = math.log((len(self.docs) - freq + 0.5) / (freq + 0.5) + 1)
        self.avgdl = num_tokens / len(self.docs)

    def score(self, query, topk):
        """
        计算文档与查询的BM25得分
        :param doc: 文档的索引
        :param query: 查询词列表
        :return: 该文档与查询的相关性得分
        """
        score_list = []
        for i in range(len(self.docs)):
            score = 0.0
            for word in query:
                if word in self.doc_freqs[i]:
                    freq = self.doc_freqs[i][word]  # 词在文档中的频率
                    # 应用BM25计算公式
                    score += (self.idf[word] * freq * (self.k1 + 1)) / (
                            freq + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avgdl))
            score_list.append(score)
        indies = [i for i, _ in heapq.nlargest(topk, enumerate(score_list), key=lambda x: x[1])]
        return [self.docs[i] for i in indies]



def main():
    SOURCE_INDEX_DATA_PATH = "/home/guweizheng/rag/basic_rag/corpus/baike_qa/baike_qa_train.json"
    text_list = []
    with open(SOURCE_INDEX_DATA_PATH, encoding="utf8") as f:
        for i, line in enumerate(f):
            if i == 20000:
                break
            text_line = json.loads(line.strip())
            text_line['content'] = text_line['title'] + ' ' + text_line['desc'] + ' ' + text_line['answer']
            text_list.append(text_line)

    bm25model = BM25Searcher(text_list)
    query = '什么人不能吃花生'
    res = bm25model.score(query, 3)
    for i in res:
        print(i)
        print('-----')
    # {'qid': 'qid_5307517376850276220', 'category': '生活-美容/塑身-减肥/塑身', 'title': '这些粉粉能不能一起吃的？我家女人想知道葛根粉木瓜粉绿茶粉苦瓜粉生 ', 'desc': '我家女人想知道\r葛根粉  木瓜粉  粉  苦瓜粉  生姜蜂蜜\r这些东西能一起吃么？', 'answer': '隔两天服1-2种，再隔两天再服另外1-2种。其中绿茶粉应单独某天服用。生姜蜂蜜和前述数种都能配手服用，数量不宜过多。', 'content': '这些粉粉能不能一起吃的？我家女人想知道葛根粉木瓜粉绿茶粉苦瓜粉生  我家女人想知道\r葛根粉  木瓜粉  粉  苦瓜粉  生姜蜂蜜\r这些东西能一起吃么？ 隔两天服1-2种，再隔两天再服另外1-2种。其中绿茶粉应单独某天服用。生姜蜂蜜和前述数种都能配手服用，数量不宜过多。'}
    # -----
    # {'qid': 'qid_8361273329149577221', 'category': '娱乐-音乐', 'title': '朴树那首<生如下花>的开头那段"吸哇啦桫咯,吸哇啦桫咯,哦哇啦噫? ', 'desc': "吸哇啦桫咯,吸哇啦桫咯,哦哇啦噫噎,吸哇拉咿呀'是哪言?什么意思?请各位达人告诉我...", 'answer': '哪一国语言都不是、什么意思都没有。他刚发行《生如夏花》这张专辑，在CCTV-3《唱片街》栏目的时候，有一个人也问了这个问题，他说他是信口捻来的。', 'content': '朴树那首<生如下花>的开头那段"吸哇啦桫咯,吸哇啦桫咯,哦哇啦噫?  吸哇啦桫咯,吸哇啦桫咯,哦哇啦噫噎,吸哇拉咿呀\'是哪言?什么意思?请各位达人告诉我... 哪一国语言都不是、什么意思都没有。他刚发行《生如夏花》这张专辑，在CCTV-3《唱片街》栏目的时候，有一个人也问了这个问题，他说他是信口捻来的。'}
    # -----
    # {'qid': 'qid_6041677890043662601', 'category': '游戏-网络游戏', 'title': '我也花了钱的，为什么总是网络连接错误？所有的服务器都试了，没为什? ', 'desc': '我也花了钱的，为什么总是连接错误？所有的服务器都试了，没为什么有防火墙，也没有路由器，为什么为什么为什么为什么', 'answer': '目前服务器已经可以正常登录了，在期间给您带来的不便，深表歉意！', 'content': '我也花了钱的，为什么总是网络连接错误？所有的服务器都试了，没为什?  我也花了钱的，为什么总是连接错误？所有的服务器都试了，没为什么有防火墙，也没有路由器，为什么为什么为什么为什么 目前服务器已经可以正常登录了，在期间给您带来的不便，深表歉意！'}
    # -----


    # seg_list = jieba.lcut_for_search(text)
    # print(seg_list)


if __name__ == '__main__':
    main()
